# cython: profile=True
from mpi4py import MPI
from numpy import dot,ones,append,array
from random import uniform
import cython
if not cython.compiled:
    from math import *
from v3dsolver.tree_manager.distribution_manager import get_domain
from v3dsolver.fluid_elements.fluid_elems import Saffman,VortexPanel
from v3dsolver.intvector import IntVec
import pdb

class Bin:
    """
    Generic bin class containing a uniform distribution of equally sized cells
    based on the global domain (not specific to any one processor).
    
    These cells are not same as cells of Tree.
    """
    def __init__(self,elems,cellsize):
        """
        Initialization of Bin class. 
        
        elems: list of elements forming bin (only local elements in case of 
               parallel implementation)
        cellsize: length of each cell
        """
        self.elems = elems
        self.domain = get_domain(elems,findglobal=True)
        self.comm = MPI.COMM_WORLD
        self.cellsize = cellsize
        self.intdomain = self.get_int_domain()
        # number of cells in x-direction (same along y and z directions too)
        self.ncells_x = long(self.intdomain[3]/self.cellsize)
        self.ncells_y = self.ncells_z = self.ncells_x
        self.hashelems = self.get_hashelems()
        self.first_elem_indices = {}
        self.next_elem_indices = {}
        self.elem_count = {}
        self.populate()
        self.assign_nbrs()

    def get_hashelems(self):
        """
        Returns a hashed representation of elems list where each element
        of the list holds information about elem type, index and processor
        rank.
        """
        hashelems = []
        for i,elem in enumerate(self.elems):
            hashelem = {'pos':elem.pos,'rank':self.comm.rank,
                        'type':type(elem),'index':i}
            hashelems.append(hashelem)
        return hashelems

    def get_int_domain(self):
        """
        Returns a modified domain such that integral number of cells come into
        it.
        """
        length = self.domain[3]
        newlength = long(length/self.cellsize)*self.cellsize + self.cellsize
        newlength += 2*self.cellsize
        newdomain = [self.domain[0]- 1.5*self.cellsize,
                     self.domain[1]-1.5*self.cellsize,
                     self.domain[2]-1.5*self.cellsize,
                     newlength]
        return newdomain
    
    def get_index(self,pos):
        """
        Returns the index of position pos which indicates what cell pos lies in.
        The returned index would be of the form of [x,y,z] where x,y and z 
        represent corresponding indices in those directions.
        """
        xmin = self.intdomain[0]
        ymin = self.intdomain[1]
        zmin = self.intdomain[2]
        x_int = (pos.x - xmin)/self.cellsize
        y_int = (pos.y - ymin)/self.cellsize
        z_int = (pos.z - zmin)/self.cellsize
        #print 'pos:', pos, 'index_b4:',(x_int,y_int,z_int) 
        x_int = long(x_int)
        y_int = long(y_int)
        z_int = long(z_int)
        #print 'pos:', pos, 'index_a5tr:',(x_int,y_int,z_int) 
        return IntVec_new(x_int,y_int,z_int)
    
    def populate(self):
        """
        Populate the bin with element indices, returns the populated 
        grid which is a 1D flattened out array corresponding to each 
        cell in 3D space and whose each element contains the index of 
        the first particle inside that cell.
        """
        # cells_first_elems would store a one dimensional list of cells
        # and would contain the index of the first elem of each cell
        size = self.ncells_x*self.ncells_y*self.ncells_z
        for i in range(len(self.elems)):
            elem = self.elems[i]
            index = self.get_index(elem.pos)
            index_1d = long(round(index.z*self.ncells_x*self.ncells_y + 
                            index.y*self.ncells_x + index.x))
            if not self.first_elem_indices.has_key(index_1d):
                self.first_elem_indices[index_1d] = i
                self.elem_count[index_1d] = 1
            else:
                current_index = self.first_elem_indices[index_1d]
                while self.next_elem_indices.has_key(current_index):
                    current_index = self.next_elem_indices[current_index]
                self.next_elem_indices[current_index] = i
                self.elem_count[index_1d] += 1
                
    def get_nbrs(self, index_1d):
        """
        Returns an array containing the 1d indices of neighbours to the 
        cell represented by index_1d. 
        """
        nx = self.ncells_x
        ny = self.ncells_y
        j = -nx*ny
        k = nx*ny
        directions = array([0,-1,1,-nx,nx,
                            -nx-1,-nx+1,nx-1,nx+1,
                            j,j-1,j+1,j-nx,j+nx,
                            j-nx-1,j-nx+1,j+nx-1,j+nx+1,
                            k,k-1,k+1,k-nx,k+nx,
                            k-nx-1,k-nx+1,k+nx-1,k+nx+1],long)
        unit_arr = ones(27,dtype=long)
        return index_1d*unit_arr + directions
    
    def assign_nbrs(self):
        """
        Returns list of neighbour indices for a cell represented by index 
        (including the cell represented by index).
        
        index is a length 3 unidimensional list whose elements contain the index 
        of a cell in x,y and z directions.
        """
        nbrlist = {}
        # creating a 2D array containing neighbour id's of each cell
        nx = self.ncells_x
        ny = self.ncells_y
        j = -nx*ny
        k = nx*ny
        directions = array([0,-1,1,-nx,nx,
                            -nx-1,-nx+1,nx-1,nx+1,
                            j,j-1,j+1,j-nx,j+nx,
                            j-nx-1,j-nx+1,j+nx-1,j+nx+1,
                            k,k-1,k+1,k-nx,k+nx,
                            k-nx-1,k-nx+1,k+nx-1,k+nx+1])
        unit_arr = ones(27)
        for i in self.first_elem_indices.keys():
            tmp_arr = i*unit_arr
            nbrlist[i] = tmp_arr + directions
        self.nbrlist = nbrlist
    
    def get_index3d(self,index_1d):
        """
        Returns the 3D index corresponding to the flattened 1d index.
        """
        index_z = index_1d/(self.ncells_x*self.ncells_y)
        index_rem = index_1d%(self.ncells_x*self.ncells_y)
        index_y = index_rem/self.ncells_x
        index_x = index_rem%self.ncells_x
        
        return (index_x, index_y,index_z)
    
    def get_lims(self,index_1d):
        """
        Returns the limits of the cell represented by index_1d.
        """
        index_z = index_1d/(self.ncells_x*self.ncells_y)
        index_rem = index_1d%(self.ncells_x*self.ncells_y)
        index_y = index_rem/self.ncells_x
        index_x = index_rem%self.ncells_x
        
        xlim = (self.intdomain[0] + index_x*self.cellsize,
                self.intdomain[0] + (index_x+1)*self.cellsize)
        ylim = (self.intdomain[1] + index_y*self.cellsize,
                self.intdomain[1] + (index_y+1)*self.cellsize)
        zlim = (self.intdomain[2] + index_z*self.cellsize,
                self.intdomain[2] + (index_z+1)*self.cellsize)
        return (xlim,ylim,zlim)
    
    def get_newelem_pos(self,index_1d):
        """
        Returns the position of a particle randomly added inside a cell
        of index index_1d.
        """
        index_z = index_1d/(self.ncells_x*self.ncells_y)
        index_rem = index_1d%(self.ncells_x*self.ncells_y)
        index_y = index_rem/self.ncells_x
        index_x = index_rem%self.ncells_x
        
        elem_x = uniform(self.intdomain[0] + index_x*self.cellsize,
                         self.intdomain[0] + (index_x+1)*self.cellsize)
        elem_y = uniform(self.intdomain[1] + index_y*self.cellsize,
                         self.intdomain[1] + (index_y+1)*self.cellsize)
        elem_z = uniform(self.intdomain[2] + index_z*self.cellsize,
                         self.intdomain[2] + (index_z+1)*self.cellsize)
        
        return Vec(elem_x,elem_y,elem_z)

class NeighbourFinder:
    """
    Generic class for finding neighbours, both serially and parallely.
    """
    def __init__(self,bin_obj,nu,dt,R):
        """
        Initialization of NeighbourFinder class.
        
        bin_obj: object of bin type
        R: the value of R taken in calculating cellsize 
        nu: kinematic viscosity
        dt: timestep
        """
        self.bin = bin_obj
        self.elems = self.bin.elems
        self.hashelems = self.bin.hashelems
        self.R = R
        self.nu = nu
        self.dt = dt
        self.cellsize = R*sqrt(self.nu*self.dt)
        comm = MPI.COMM_WORLD
        self.cellelems_list = {}
        reqlist = self.send_particles(comm)
        self.remote_nbrs = self.recv_particles(comm)
        MPI.Prequest.Waitall(reqlist)
        
    def are_neighbours(self,base_elem,hashnbr):
        """
        Here, hashnbr is the dict based hashed representation of nbr.
        Returns True if following conditions are satisfied:
        1. hashnbr['type']==Saffman
        2. nbr_elem is within `cellsize` distance from base_elem
        """
        if not hashnbr['type']==Saffman:
            return False
        else:
            pos1 = base_elem.pos
            pos2 = hashnbr['pos']
            rvec = pos1 - pos2
            r = rvec.length()
            if round(r/self.cellsize,5)<=1:
                return True
            else:
                return False
    
    def get_neighbours_direct(self,element):
        """
        Return particles in the neighbourhood (within radius R*sqrt(nu*dt))
        of elem. Uses direct linear search over all the 
        particles.
        """
        nbrs = []
        for hashnbr in self.hashelems:
            if self.are_neighbours(element,hashnbr):
                nbrs.append(hashnbr)
        return nbrs
        
    def get_local_searchspace(self,element):
        """
        Returns the local searchspace containing all elements from 
        neighboring 27 cells.
        """
        pos = element.pos
        parentcell = self.bin.get_index(pos)
        index_1d = long(round((parentcell.z*self.bin.ncells_x*self.bin.ncells_y + 
                    parentcell.y*self.bin.ncells_x + parentcell.x)))
        
        nbrlist = self.bin.nbrlist[index_1d]
        nbrlist = [long(round(nbr_cell)) for nbr_cell in nbrlist]
        
        nbrs = []
        for nbr_cell in nbrlist:
            self._update_elems(nbr_cell)
            if self.cellelems_list.has_key(nbr_cell):
                nbrs += self.cellelems_list[nbr_cell]
        return nbrs
        
    def _update_elems(self,index_1d):
        """
        If self.cellelems_list[index_1d] is empty, it will fill it with
        elements of that cell (in the format [hashkey,elempos]), 
        otherwise, keep intact.
        """
        if (self.bin.first_elem_indices.has_key(index_1d) and 
            not self.cellelems_list.has_key(index_1d)):
            cellelems = []
            current_index = self.bin.first_elem_indices[index_1d]
            while True:
                current_hashelem = self.hashelems[current_index]
                cellelems.append(current_hashelem)
                if not self.bin.next_elem_indices.has_key(current_index):
                    break
                current_index = self.bin.next_elem_indices[current_index]
            self.cellelems_list[index_1d] = cellelems
            
    def get_neighbours_serial(self,element):
        """
        Return particles in the neighbourhood (within radius R*sqrt(nu*dt))
        of elem, using serial bin data structure.
        """
        pos = element.pos
        searchspace = self.get_local_searchspace(element)
        nbrs = []
        for hashnbr in searchspace:
            if self.are_neighbours(element,hashnbr):
                nbrs.append(hashnbr)
        return nbrs

    def send_particles(self,comm):
        """
        Traverses through own bin and transfers particles of own cell to 
        appropriate processors.
        """
        proc_map = comm.allgather(self.bin.first_elem_indices)
        send_data = {}
        for i in range(comm.size):
            if not i==comm.rank:
                ith_proc_map = proc_map[i]
                send_data.update({i:[]})
                for j in (self.bin.first_elem_indices).keys():
                        for nbr in self.bin.nbrlist[j]:
                            if ith_proc_map.has_key(nbr):
                                self._update_elems(j)
                                send_data[i].append([nbr,
                                self.cellelems_list[j]])

        reqlist = []
        for i in range(comm.size):
            if not i==comm.rank:
                reqlist += [comm.isend(send_data[i],dest=i,tag=6)]
                print 'sent by', i
        return reqlist
    
    def recv_particles(self,comm):
        """
        Receive particles from all (except self) processors
        """
        remote_nbrs = {}
        for i in (self.bin.first_elem_indices).keys():
            remote_nbrs.update({i:[]})
        for i in range(comm.size):
            if not i==comm.rank:
                recvd_data = comm.recv(source=i,tag=6)
                print 'recvd by', i
                for celldata in recvd_data:
                    cellid = celldata[0]
                    remote_nbrs[cellid] = celldata[1]
        return remote_nbrs
    
    def get_neighbours_parallel(self,element,comm=MPI.COMM_WORLD):
        """
        Return particles in the neighbourhood (within radius R*sqrt(nu*dt))
        of element denoted by elem, using parallel bin data structure.
        
        comm: MPI communicator object
        
        Methodology for parallel neighbour finding:
        1. [Client] Traverses own bin for marking particles in neighbouring cells 
           as possible neighbours (done in Constructor)
        2. [Server] Sends different processors particles which can act as their
           neighbours (done in Constructor)
        3. [Client] Based on received non-local neighbour particles, it updates
           possible neighbour data. (done in Constructor)
        4. [Client] Conducts a search in the obtained search space of possible 
           neighbours, it returns particles which are TRUE neighbours to given
           particle. (done in this function)
        """
        local_nbrs = self.get_local_searchspace(element)
        pos = element.pos
        parentcell = self.bin.get_index(pos)
        index_1d = (parentcell.z*self.bin.ncells_x*self.bin.ncells_y + 
                    parentcell.y*self.bin.ncells_x + parentcell.x)
        remote_nbrs = self.remote_nbrs[index_1d]
        search_space = local_nbrs + remote_nbrs
        nbrs = []
        for hashnbr in search_space:
            if self.are_neighbours(element,hashnbr):
                nbrs.append(hashnbr)
        return nbrs
